import torch
import torch.nn.functional as F



def mse_between_patchs_loss(
    predict_img, label_img, *args
):

    # diff = (predict_img-label_img)
    # loss = torch.mean(diff * diff)

    loss = F.mse_loss(predict_img, label_img)

    return loss


def DG_loss(
    predict_img, label_img, noisy_img
):

    loss = F.mse_loss

    loss_denominator = loss(predict_img, label_img)
    loss_numerator = loss(noisy_img, label_img)
    # loss_numerator = 1

    # Calculate DG Loss
    dg_loss = torch.log10(loss_denominator / loss_numerator)

    return dg_loss


def TV_loss_formula(predict_img, *args):

    factor = predict_img.shape[0] * predict_img.shape[1]

    diff_along_width = predict_img[:, :, :-1, :-1] - predict_img[:, :, :-1, 1:]
    diff_along_height = predict_img[:, :, :-1, :-1] - predict_img[:, :, 1:, :-1]
    eps = 1e-7

    loss = (diff_along_width*diff_along_width
        + diff_along_height*diff_along_height+eps).sqrt().sum() /factor

    return loss


def TV_loss_tf(predict_img, *args):

    factor = predict_img.shape[0] * predict_img.shape[1]

    diff_along_width = predict_img[:, :, :, :-1] - predict_img[:, :, :, 1:]
    diff_along_height = predict_img[:, :, :-1, :] - predict_img[:, :, 1:, :]

    loss = (diff_along_width.abs().sum()
        + diff_along_height.abs().sum()) / factor

    return loss


def TV_loss_CAM(predict_img, *args):

    x = predict_img

    # Initialize Variables
    batch_size = x.size()[0]
    batch_size = x.size()[0]
    h_x = x.size()[2]
    w_x = x.size()[3]

    count_h = tensor_size(x[:, :, 1:, :])
    count_w = tensor_size(x[:, :, :, 1:])

    h_tv = torch.pow((x[:, :, 1:, :] - x[:, :, :h_x - 1, :]), 2).sum()
    w_tv = torch.pow((x[:, :, :, 1:] - x[:, :, :, :w_x - 1]), 2).sum()

    loss = 2 * (h_tv / count_h + w_tv / count_w) / batch_size

    return loss


def tensor_size(t) :
    return t.size()[1] * t.size()[2] * t.size()[3]

